#!/usr/bin/env python3
"""
现代化网页版聊天应用
基于Flask框架，提供现代化的Web界面
"""

from flask import Flask, render_template, request, jsonify, session
from flask_socketio import SocketIO
import json
import os
import time
import threading
from datetime import datetime
from typing import Dict, Any, List

# 导入现有的核心模块
from models.factory import ModelFactory
from utils.config_manager import ConfigManager
from utils.chat_history import ChatHistoryManager
from utils.knowledge_base import KnowledgeBase
from utils.context_manager import context_manager

app = Flask(__name__, 
           template_folder='web_templates',
           static_folder='web_static')
app.secret_key = 'your-secret-key-here'  # 生产环境请更换
socketio = SocketIO(app, cors_allowed_origins="*")

# 初始化核心组件
model_factory = ModelFactory()
config_manager = ConfigManager("config.json")
chat_history = ChatHistoryManager("chat_history.json")
knowledge_base = KnowledgeBase(dir_path="knowledge_base")

# 获取模型配置
models = model_factory.get_model_configs()

# 从配置中加载API密钥
for model_name in models:
    api_key = config_manager.get_api_key(model_name)
    if api_key:
        models[model_name]["api_key"] = api_key

# 当前会话状态
sessions: Dict[str, Dict[str, Any]] = {}

class WebChatSession:
    """Web聊天会话管理"""
    
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.messages: List[Dict[str, str]] = []
        self.current_model = list(models.keys())[0]
        self.enable_thinking = False
        self.current_role = config_manager.get_current_role()
        
        # 加载角色配置
        self.roles = config_manager.get_roles()
        
        # 初始化会话
        sessions[session_id] = {
            'messages': self.messages,
            'current_model': self.current_model,
            'enable_thinking': self.enable_thinking,
            'current_role': self.current_role
        }
    
    def add_message(self, role: str, content: str):
        """添加消息到会话"""
        message = {
            'role': role,
            'content': content,
            'timestamp': datetime.now().isoformat(),
            'id': f"msg_{len(self.messages)}_{int(time.time())}"
        }
        self.messages.append(message)
        
        # 保存到历史记录
        if role in ['user', 'assistant']:
            chat_history.add_message(role, content)
        
        return message
    
    def get_recent_messages(self, limit: int = 10) -> List[Dict[str, str]]:
        """获取最近的聊天消息"""
        return self.messages[-limit:]
    
    def clear_messages(self):
        """清空聊天消息"""
        self.messages.clear()
        chat_history.clear_current_session()

@app.route('/')
def index():
    """主页面"""
    session_id = request.args.get('session_id') or f'session_{int(time.time())}'
    
    if session_id not in sessions:
        WebChatSession(session_id)
    
    # 获取所有模型和角色
    model_names = list(models.keys())
    role_names = list(config_manager.get_roles().keys())
    
    return render_template('index.html', 
                          session_id=session_id,
                          models=model_names,
                          roles=role_names,
                          current_model=sessions[session_id]['current_model'],
                          current_role=sessions[session_id]['current_role'])

@app.route('/api/chat', methods=['POST'])
def chat_api():
    """聊天API接口"""
    try:
        data = request.get_json()
        session_id = data.get('session_id')
        message = data.get('message')
        
        if not session_id or session_id not in sessions:
            return jsonify({'error': 'Invalid session'}), 400
        
        if not message:
            return jsonify({'error': 'Message is required'}), 400
        
        session_data = sessions[session_id]
        
        # 添加用户消息
        user_message = session_data['messages'].append({
            'role': 'user',
            'content': message,
            'timestamp': datetime.now().isoformat()
        })
        
        # 获取AI响应（异步）
        def get_ai_response():
            try:
                # 获取当前模型
                current_model = session_data['current_model']
                model_config = models[current_model]
                
                # 创建模型实例
                model = model_factory.create_model(current_model, model_config)
                
                # 获取角色配置
                current_role = session_data['current_role']
                role_config = config_manager.get_role_config(current_role)
                system_prompt = role_config.get('system_prompt', '你是一个友善的AI助手')
                
                # 构建消息历史
                messages = [
                    {'role': 'system', 'content': system_prompt}
                ]
                
                # 添加上下文消息
                recent_messages = context_manager.get_recent_context(
                    session_data['messages'], look_back=5
                )
                messages.extend(recent_messages)
                
                # 添加当前消息
                messages.append({'role': 'user', 'content': message})
                
                # 获取响应
                response = model.generate_response(
                    messages=messages,
                    max_tokens=config_manager.get_setting('max_tokens', 2000),
                    temperature=config_manager.get_setting('temperature', 0.7),
                    enable_thinking=session_data['enable_thinking']
                )
                
                # 添加助手响应
                assistant_message = {
                    'role': 'assistant',
                    'content': response,
                    'timestamp': datetime.now().isoformat(),
                    'model': current_model
                }
                session_data['messages'].append(assistant_message)
                
                # 发送到前端
                socketio.emit('chat_response', {
                    'session_id': session_id,
                    'message': assistant_message
                })
                
            except Exception as e:
                error_msg = f"Error: {str(e)}"
                socketio.emit('chat_error', {
                    'session_id': session_id,
                    'error': error_msg
                })
        
        # 在新线程中处理AI响应
        threading.Thread(target=get_ai_response, daemon=True).start()
        
        return jsonify({'status': 'processing'})
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/settings', methods=['GET', 'POST'])
def settings_api():
    """设置API接口"""
    if request.method == 'GET':
        # 获取当前设置
        session_id = request.args.get('session_id')
        if session_id and session_id in sessions:
            session_data = sessions[session_id]
            return jsonify({
                'current_model': session_data['current_model'],
                'enable_thinking': session_data['enable_thinking'],
                'current_role': session_data['current_role'],
                'typing_speed': config_manager.get_typing_speed()
            })
        return jsonify({'error': 'Session not found'}), 404
    
    else:  # POST
        data = request.get_json()
        session_id = data.get('session_id')
        
        if not session_id or session_id not in sessions:
            return jsonify({'error': 'Invalid session'}), 400
        
        session_data = sessions[session_id]
        
        # 更新设置
        if 'current_model' in data:
            session_data['current_model'] = data['current_model']
            config_manager.save_current_model(data['current_model'])
        
        if 'enable_thinking' in data:
            session_data['enable_thinking'] = data['enable_thinking']
        
        if 'current_role' in data:
            session_data['current_role'] = data['current_role']
            config_manager.save_current_role(data['current_role'])
        
        if 'typing_speed' in data:
            config_manager.save_typing_speed(data['typing_speed'])
        
        return jsonify({'status': 'success'})

@app.route('/api/history')
def chat_history_api():
    """获取聊天历史"""
    session_id = request.args.get('session_id')
    
    if not session_id or session_id not in sessions:
        return jsonify({'error': 'Invalid session'}), 400
    
    session_data = sessions[session_id]
    return jsonify({'messages': session_data['messages']})

@app.route('/api/clear', methods=['POST'])
def clear_chat_api():
    """清空聊天"""
    data = request.get_json()
    session_id = data.get('session_id')
    
    if not session_id or session_id not in sessions:
        return jsonify({'error': 'Invalid session'}), 400
    
    sessions[session_id]['messages'].clear()
    chat_history.clear_current_session()
    
    return jsonify({'status': 'success'})

@socketio.on('connect')
def handle_connect():
    print('Client connected')

@socketio.on('disconnect')
def handle_disconnect():
    print('Client disconnected')

if __name__ == '__main__':
    # 创建模板和静态文件目录
    os.makedirs('web_templates', exist_ok=True)
    os.makedirs('web_static', exist_ok=True)
    os.makedirs('web_static/css', exist_ok=True)
    os.makedirs('web_static/js', exist_ok=True)
    
    # 运行应用
    print("🚀 启动现代化网页版聊天应用...")
    print("📱 访问地址: http://localhost:5000")
    socketio.run(app, host='0.0.0.0', port=5000, debug=True)